"""
Shared base functionality for server testing.

This module contains the common setup, cleanup, and utility functions
used by both server.py and server_llamacpp.py tests.
"""

import unittest
import subprocess
import psutil
import asyncio
import socket
import time
from threading import Thread
import sys
import io
import httpx
import argparse
import contextlib
from unittest.mock import patch
import urllib.request
import os
import requests

try:
    from openai import OpenAI, AsyncOpenAI
except ImportError as e:
    raise ImportError("You must `pip install openai` to run this test", e)

# Import huggingface_hub for patching in offline mode
try:
    from huggingface_hub import snapshot_download as original_snapshot_download
except ImportError:
    # If huggingface_hub is not installed, create a dummy function
    def original_snapshot_download(*args, **kwargs):
        raise ImportError("huggingface_hub is not installed")


MODEL_NAME = "Qwen2.5-0.5B-Instruct-CPU"
# This list must include all models that could be accessed in offline testing
MODELS_UNDER_TEST = [
    MODEL_NAME,
    "Llama-3.2-1B-Instruct-CPU",  # used in test_001_test_simultaneous_load_requests
]
MODEL_CHECKPOINT = "amd/Qwen2.5-0.5B-Instruct-quantized_int4-float16-cpu-onnx"
PORT = 8000


# Global variable for server binary (can be overridden via --server-binary)
# Default to finding lemonade-server-dev in the same venv as the running Python
def _get_default_server_binary():
    """Get the default server binary path based on the current Python interpreter."""
    import platform

    # Get the directory containing the Python executable
    # For venv: .venv/Scripts/python.exe (Windows) or .venv/bin/python (Linux/Mac)
    python_dir = os.path.dirname(sys.executable)
    if platform.system() == "Windows":
        return os.path.join(python_dir, "lemonade-server-dev.exe")
    else:
        return os.path.join(python_dir, "lemonade-server-dev")


SERVER_BINARY = _get_default_server_binary()


def is_cpp_server():
    """Check if we're testing the C++ server instead of Python.

    Returns True if --server-binary argument was provided (i.e., not using the default Python server).
    """
    # If --server-binary was provided and it's not our default venv binary, we're testing a custom binary (C++ server)
    return SERVER_BINARY != _get_default_server_binary()


def stop_lemonade():
    """
    Kill the lemonade server and stop the model
    """
    # Kill the server subprocess
    print("\n=== Stopping Lemonade ===")

    result = subprocess.run(
        [SERVER_BINARY, "stop"],
        capture_output=True,
        text=True,
    )
    print(result.stdout)


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description="Test lemonade server", add_help=False)
    parser.add_argument(
        "--offline", action="store_true", help="Run tests in offline mode"
    )
    parser.add_argument(
        "--server-binary",
        type=str,
        default=_get_default_server_binary(),
        help="Path to server binary (default: lemonade-server-dev in venv)",
    )
    # Use parse_known_args to ignore unknown arguments (like positional 'backend' in server_llamacpp.py)
    args, unknown = parser.parse_known_args()

    # Update global SERVER_BINARY
    global SERVER_BINARY
    SERVER_BINARY = args.server_binary

    return args


@contextlib.contextmanager
def simulate_offline_mode():
    """
    Context manager that simulates a fully offline environment except
    for local connections needed for testing.

    This patches multiple network-related functions to prevent any
    external network access during tests.
    """
    original_create_connection = socket.create_connection

    def mock_create_connection(address, *args, **kwargs):
        host, port = address
        # Allow connections to localhost for testing
        if host == "localhost" or host == "127.0.0.1":
            return original_create_connection(address, *args, **kwargs)
        # Block all other connections
        raise socket.error("Network access disabled for offline testing")

    # Define a function that raises an error for non-local requests
    def block_external_requests(original_func):
        def wrapper(url, *args, **kwargs):
            # Allow localhost requests
            if url.startswith(
                (
                    "http://localhost",
                    "https://localhost",
                    "http://127.0.0.1",
                    "https://127.0.0.1",
                )
            ):
                return original_func(url, *args, **kwargs)
            raise ConnectionError(f"Offline mode: network request blocked to {url}")

        return wrapper

    # Apply all necessary patches to simulate offline mode
    with patch("socket.create_connection", side_effect=mock_create_connection):
        with patch(
            "huggingface_hub.snapshot_download",
            side_effect=lambda *args, **kwargs: (
                kwargs.get("local_files_only", False)
                and original_snapshot_download(*args, **kwargs)
                or (_ for _ in ()).throw(
                    ValueError("Offline mode: network connection attempted")
                )
            ),
        ):
            # Also patch urllib and requests to block external requests
            with patch(
                "urllib.request.urlopen",
                side_effect=block_external_requests(urllib.request.urlopen),
            ):
                with patch(
                    "http.client.HTTPConnection.connect",
                    side_effect=lambda self, *args, **kwargs: (
                        None
                        if self.host in ("localhost", "127.0.0.1")
                        else (_ for _ in ()).throw(
                            ConnectionError("Offline mode: connection blocked")
                        )
                    ),
                ):
                    # Set environment variable to signal offline mode
                    os.environ["LEMONADE_OFFLINE_TEST"] = "1"
                    try:
                        yield
                    finally:
                        # Clean up environment variable
                        if "LEMONADE_OFFLINE_TEST" in os.environ:
                            del os.environ["LEMONADE_OFFLINE_TEST"]


def ensure_model_is_cached():
    """
    Make sure the test model is downloaded and cached locally before running in offline mode.
    """
    try:
        # Call server binary pull to download the model
        for model_name in MODELS_UNDER_TEST:
            subprocess.run(
                [SERVER_BINARY, "pull", model_name],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
            )
            print(f"Model {model_name} successfully pulled and available in cache")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Failed to download model: {e}")
        return False


class ServerTestingBase(unittest.IsolatedAsyncioTestCase):
    """Base class containing only shared setup/cleanup functionality, no test methods."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Allow subclasses to set the llamacpp backend
        self.llamacpp_backend = getattr(self, "llamacpp_backend", None)

    def setUp(self):
        """
        Start lemonade server process
        """
        print("\n=== Starting new test ===")
        self.base_url = f"http://localhost:{PORT}/api/v1"
        self.messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Who won the world series in 2020?"},
            {"role": "assistant", "content": "The LA Dodgers won in 2020."},
            {"role": "user", "content": "What was the best play?"},
        ]

        # Ensure we stop lemonade
        stop_lemonade()

        # Build the command to start the server
        cmd = [SERVER_BINARY, "serve"]

        # Add --no-tray option on Windows
        if os.name == "nt":
            cmd.append("--no-tray")

        # Add debug logging for CI environments
        cmd.extend(["--log-level", "debug"])

        # Add llamacpp backend option if specified
        if self.llamacpp_backend:
            cmd.extend(["--llamacpp", self.llamacpp_backend])

        # Add any additional server arguments
        if hasattr(self.__class__, "additional_server_args"):
            cmd.extend(self.__class__.additional_server_args)

        # Start the lemonade server
        lemonade_process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
            encoding="utf-8",
            errors="replace",  # Replace any non-UTF-8 characters to prevent crashes
            env=os.environ.copy(),
        )

        # Print stdout and stderr in real-time
        # Use separate threads for stdout and stderr to prevent blocking
        def print_stdout():
            for line in lemonade_process.stdout:
                print(f"[stdout] {line.strip()}")

        def print_stderr():
            for line in lemonade_process.stderr:
                print(f"[stderr] {line.strip()}")

        # Start output threads
        stdout_thread = Thread(target=print_stdout, daemon=True)
        stderr_thread = Thread(target=print_stderr, daemon=True)
        stdout_thread.start()
        stderr_thread.start()

        # Wait for the server to start by checking port 8000
        start_time = time.time()
        while True:
            if time.time() - start_time > 60:
                raise TimeoutError("Server failed to start within 60 seconds")
            try:
                conn = socket.create_connection(("localhost", PORT))
                conn.close()
                break
            except socket.error:
                time.sleep(1)

        # Wait a few other seconds after the port is available
        time.sleep(5)

        print("Server started successfully")

        self.addCleanup(stop_lemonade)

        # Ensure stdout can handle Unicode
        if sys.stdout.encoding != "utf-8":
            sys.stdout = io.TextIOWrapper(
                sys.stdout.buffer, encoding="utf-8", errors="replace"
            )
            sys.stderr = io.TextIOWrapper(
                sys.stderr.buffer, encoding="utf-8", errors="replace"
            )


def run_server_tests_with_class(
    test_class, description="SERVER TESTS", offline=None, additional_args=None
):
    """Utility function to run server tests with a given test class.

    Args:
        test_class: The unittest.TestCase class to run
        description: Description for the test run
        offline: Whether to run in offline mode (defaults to parsed --offline arg)
        additional_args: List of additional command-line arguments to pass to the server
    """
    # Always parse args to set SERVER_BINARY global
    args = parse_args()

    # If offline parameter is not provided, use parsed value
    if offline is None:
        offline = args.offline

    # Store additional args in a class variable so setUp can access them
    if additional_args:
        test_class.additional_server_args = additional_args

    if offline:
        print(f"\n=== STARTING {description} IN OFFLINE MODE ===")

        if not ensure_model_is_cached():
            print("ERROR: Unable to cache the model needed for offline testing")
            sys.exit(1)

        print("Model is cached. Running tests with network access disabled...")

        # Create a new test suite
        test_loader = unittest.TestLoader()
        test_suite = test_loader.loadTestsFromTestCase(test_class)

        # Run the tests in offline mode
        with simulate_offline_mode():
            result = unittest.TextTestRunner(
                verbosity=2, buffer=False, failfast=True
            ).run(test_suite)

        # Set exit code based on test results
        sys.exit(0 if result.wasSuccessful() else 1)
    else:
        print(f"\n=== STARTING {description} IN NORMAL MODE ===")
        # Create a new test suite for the specific class
        test_loader = unittest.TestLoader()
        test_suite = test_loader.loadTestsFromTestCase(test_class)
        # Use verbosity=2 to show test names, buffer=False to see output in real-time,
        # and failfast=True to stop on first failure and show the error immediately
        result = unittest.TextTestRunner(verbosity=2, buffer=False, failfast=True).run(
            test_suite
        )

        # Set exit code based on test results
        sys.exit(0 if result.wasSuccessful() else 1)


# This file was originally licensed under Apache 2.0. It has been modified.
# Modifications Copyright (c) 2025 AMD
